#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 09:55:30
# @LastEditTime: 2025-03-14 09:55:35
# @Description: AWS监控核心模块
###

import logging
from typing import Dict, List
import boto3
from botocore.exceptions import ClientError

class AWSMonitor:
    """AWS监控基类"""
    
    def __init__(self, config: Dict):
        """初始化
        
        Args:
            config: 配置信息，包含ak和region
        """
        self.config = config
        self.session = boto3.Session(
            aws_access_key_id=config['access_key'],
            aws_secret_access_key=config['secret_key'],
            region_name=config['region']
        )
        self.cloudwatch = self.session.client('cloudwatch')
        
    def get_metric_data(self, namespace: str, metric_name: str, 
                       dimensions: List[Dict], period: int = 300) -> Dict:
        """获取监控数据
        
        Args:
            namespace: 命名空间
            metric_name: 指标名称
            dimensions: 维度列表
            period: 统计周期（秒）
        """
        try:
            response = self.cloudwatch.get_metric_data(
                MetricDataQueries=[{
                    'Id': 'm1',
                    'MetricStat': {
                        'Metric': {
                            'Namespace': namespace,
                            'MetricName': metric_name,
                            'Dimensions': dimensions
                        },
                        'Period': period,
                        'Stat': 'Average'
                    }
                }],
                StartTime='-PT1H',  # 最近1小时
                EndTime='PT0H'
            )
            return response
        except ClientError as e:
            logging.error(f"获取监控数据失败: {e}")
            return {}
            
    def create_alarm(self, alarm_name: str, namespace: str,
                    metric_name: str, dimensions: List[Dict],
                    threshold: float, period: int = 300) -> None:
        """创建告警
        
        Args:
            alarm_name: 告警名称
            namespace: 命名空间
            metric_name: 指标名称
            dimensions: 维度列表
            threshold: 阈值
            period: 统计周期（秒）
        """
        try:
            self.cloudwatch.put_metric_alarm(
                AlarmName=alarm_name,
                MetricName=metric_name,
                Namespace=namespace,
                Dimensions=dimensions,
                Period=period,
                EvaluationPeriods=3,
                Threshold=threshold,
                ComparisonOperator='GreaterThanThreshold',
                Statistic='Average'
            )
            logging.info(f"创建告警成功: {alarm_name}")
        except ClientError as e:
            logging.error(f"创建告警失败: {e}")
            
    def list_alarms(self) -> List[Dict]:
        """获取告警列表"""
        try:
            response = self.cloudwatch.describe_alarms()
            return response.get('MetricAlarms', [])
        except ClientError as e:
            logging.error(f"获取告警列表失败: {e}")
            return []
            
    def get_alarm_history(self, alarm_name: str) -> List[Dict]:
        """获取告警历史
        
        Args:
            alarm_name: 告警名称
        """
        try:
            response = self.cloudwatch.describe_alarm_history(
                AlarmName=alarm_name,
                HistoryItemType='StateUpdate'
            )
            return response.get('AlarmHistoryItems', [])
        except ClientError as e:
            logging.error(f"获取告警历史失败: {e}")
            return []
            
    def aggregate_alarms(self, alarms: List[Dict]) -> Dict:
        """聚合告警信息
        
        Args:
            alarms: 告警列表
        """
        result = {
            'total': len(alarms),
            'in_alarm': 0,
            'insufficient_data': 0,
            'ok': 0,
            'top_instances': []
        }
        
        # 统计各状态告警数量
        for alarm in alarms:
            state = alarm.get('StateValue', '')
            if state == 'ALARM':
                result['in_alarm'] += 1
            elif state == 'INSUFFICIENT_DATA':
                result['insufficient_data'] += 1
            elif state == 'OK':
                result['ok'] += 1
                
            # 收集异常实例信息
            if state == 'ALARM':
                instance_info = {
                    'name': alarm.get('AlarmName'),
                    'metric': alarm.get('MetricName'),
                    'namespace': alarm.get('Namespace'),
                    'value': alarm.get('StateReason'),
                    'time': alarm.get('StateUpdatedTimestamp')
                }
                result['top_instances'].append(instance_info)
                
        # 只保留最近的10个告警实例
        result['top_instances'] = sorted(
            result['top_instances'],
            key=lambda x: x['time'],
            reverse=True
        )[:10]
        
        return result
